import os
import os.path as osp
import sys


from app.pdf_extract_kit.tasks.base_task import TaskFlow
import app.pdf_extract_kit.tasks # 勿动！！！！
from app.pdf_extract_kit.utils.layout_utils import ocr_layout_merge
from app.project.doc_to_recommendation.utils.log_utils import get_logger

sys.path.append(osp.join(os.path.dirname(os.path.abspath(__file__)), '..'))
from app.project.doc_to_recommendation.utils.config_loader import load_config, get_config_path, initialize_tasks_and_models



TASK_NAME_OCR = 'ocr'
TASK_NAME_LAYOUT = 'layout_detection'
CONFIG_PATH_DEFAULT = get_config_path("ocr_layout.yaml")

def process(file_path=None, save_dir=None, config_path=None):
    log = get_logger(__name__)
    log.info("读配置")
    config = load_config(CONFIG_PATH_DEFAULT) if config_path is None else load_config(config_path)

    return do_process(config=config, file_path=file_path, save_dir=save_dir, log=log)

def do_process(config: dict, file_path=None, save_dir=None, log=None):
    if log is None:
        log = get_logger(__name__)

    log.info("初始化模型")
    task_instances = initialize_tasks_and_models(config)
    # get input and output path from config
    input_data = file_path if file_path is not None else config.get('inputs', None)
    file_name = os.path.splitext(os.path.basename(input_data))[0]
    visualize_save_path = config.get('log_outputs', 'logs/OCR_LAYOUT') + '/' + file_name
    result_path_img = os.path.join(save_dir, file_name) if save_dir is not None else os.path.join('outputs', file_name)
    img_save = config['img_process'].get('img_save', False)
    log.info("任务流开始")
    cfg = {}
    cfg.update(config['tasks'])
    # 添加pptx转pdf配置
    cfg.update({"ppt_load": config.get('ppt_load', {})})
    flow = TaskFlow([task_instances[TASK_NAME_OCR], task_instances[TASK_NAME_LAYOUT]], cfg)
    results, images = flow.process(input_data, visualize_save_path)

    # print(f'Task done, results can be found at {config.get('outputs')}')
    return ocr_layout_extract(input_data, images, results[TASK_NAME_OCR], results[TASK_NAME_LAYOUT], result_path_img, img_save=img_save)

def ocr_layout_extract(fpath, images, ocr_results, layout_results, save_dir="", img_save=False):
    os.makedirs(save_dir, exist_ok=True) if img_save else None
    basename = os.path.basename(fpath)[:-4]
    figures = []
    contents = []
    log = get_logger("ocr_layout_extract")
    log.info("ocr_layout_extract start")
    for idx, (image, ocr_result, layout_result) in enumerate(zip(images, ocr_results, layout_results)):
        ol_results = ocr_layout_merge(ocr_result, layout_result.__dict__['boxes'].xyxy, layout_result.__dict__['boxes'].cls, layout_result.__dict__['boxes'].conf)

        if len(ol_results['image_boxes']):
            figure_temp = []
            for i, image_box in enumerate(ol_results['image_boxes']):
                file_name = os.path.join(save_dir, f"{basename}-page{idx}-no{i}.png")
                img_cut = image.crop(image_box.xyxy)
                figure_temp.append({"image": img_cut, "name": file_name, 'label': image_box.text})
                img_cut.save(file_name) if img_save else None
            figures.append(figure_temp)
        if len(ol_results['text_boxes']):
            contents.append([box.text for box in ol_results['text_boxes']])
    log.info("ocr_layout_extract end")
    return contents, figures

# def visualize_ocr_results(fpath, ocr_results, save_dir=""):
#     def visualize_image(image, ocr_res, save_path="", cate2color={}):
#         """plot each result's bbox and category on image.
#
#         Args:
#             image: PIL.Image.Image
#             ocr_res: list of ocr det and rec, whose format following the results of self.predict_image function
#             save_path: path to save visualized image
#         """
#         draw = ImageDraw.Draw(image)
#         for res in ocr_res:
#             box_color = cate2color.get(res['category_type'], (0, 255, 0))
#             x_min, y_min = int(res['bbox'][0]), int(res['bbox'][1])
#             x_max, y_max = int(res['bbox'][2]), int(res['bbox'][3])
#             draw.rectangle([x_min, y_min, x_max, y_max], fill=None, outline=box_color, width=1)
#             draw.text((x_min, y_min), res['category_type'], (255, 0, 0))
#         if save_path:
#             image.save(save_path)
#
#     if fpath.lower().endswith(('.pdf')):
#         images = load_pdf(fpath)
#     else:
#         images = load_ppt(fpath)
#     basename = os.path.basename(fpath)[:-4]
#     for idx, (ocr_result, image) in enumerate(zip(ocr_results, images)):
#         visualize_image(image, ocr_result, os.path.join(save_dir, basename, f"page_cluster_{idx+1}.jpg"))

if __name__ == "__main__":
    # args = parse_args()
    # main(args.config)
    a, b = process(f"/home/XXX/Documents/prod_dc/prdt_doc_to_rec_manual-main/企业知识管理系统产品推介手册.pdf")
    print(a, b)